#!/usr/bin/python
#
# Name: rtc_mix
# Date: 17-JAN-2008
# Author: Mike Witthoeft
#
# Description:
# Extract mixing percentages from an AUTOSTRUCTURE olg file.  There
# are three modes of operation: (1) full, show mixing percentages for
# individual levels, (2) conf, sum percentages over configuration
# labels, and (3) term, sum percentages over term labels.
#
# Options:
#  -h, --help            show this help message and exit
#  -m MODE, --mode=MODE  select operation mode: full (default), conf, term
#  -t TOL, --tol=TOL     specify number of decimal points to print
#  -f FILE, --filter=FILE
#                        file containing level indices
#  -o FILE, --output=FILE
#                        write to file instead of screen
#
from rap_core import *
from optparse import OptionParser
import os.path
import sys
import re

def sum_config(levels,coeffs):
    """ 
    Return dictionary where the square of the mixing coefficients are summed 
    across configuration labels.
    """
    out={}
    for (k,v) in coeffs.iteritems():
        c=levels.configuration(k)
        if c in out.keys():
            out[c]+=v**2
        else:
            out[c]=v**2
    return out

def sum_term(levels,coeffs):
    """ 
    Return dictionary where the square of the mixing coefficients are summed 
    across term labels.
    """
    out={}
    for (k,v) in coeffs.iteritems():
        c=levels.term(k)
        if c in out.keys():
            out[c]+=v**2
        else:
            out[c]=v**2
    return out

def sum_full(levels,coeffs):
    """
    Prepare dictionary with all mixing percentages.
    """
    out={}
    keys=coeffs.keys()
    maxlen=0
    for k in keys: maxlen=max(maxlen,len(str(k)))
    iform='%-'+str(maxlen+1)+'s'
    for (k,v) in coeffs.iteritems():
        c=iform%(str(k)+':')+levels.label(k)
        out[c]=v**2
    return out

def show_coeffs(cout,label,tol=2,filename=None):
    """ Print mixing percentages to screen. """
    # check tolerance
    try:
        ttol=int(tol)
    except:
        ttol=2
    else:
        if ttol < 1:
            ttol=1
        elif ttol > 6:
            ttol=6

    # get format strings
    items=cout.keys()
    maxlen=0
    for it in items:
        tit=str(it)
        maxlen=max(maxlen,len(tit))
    aform='%-'+str(maxlen)+'s'
    cform='%'+str(ttol+2)+'.'+str(ttol)+'f'
    clform='%-'+str(ttol+2)+'s'

    # sort coefficients
    out=[]
    for it in items: out.append((cout[it],it))
    out.sort()
    out.reverse()

    # print to screen or file
    if filename:
        try:
            fout=open(filename,'a')
        except:
            return False
        fout.write(aform%str(label)+'   '+clform%'%'+'\n')
        for (p,i) in out: 
            if p <= 0.0: continue
            if log10(p) >= -ttol: fout.write(aform%i+'   '+cform%p+'\n')
        fout.close()
    else:
        print aform%label+'   '+clform%'frac'
        for (p,i) in out: 
            if p <= 0.0: continue
            if log10(p) >= -ttol: print aform%i+'   '+cform%p
    return True

def read_levlist(filename):
    """ Read level list from filename. """
    try:
        fin=open(filename,'r')
    except:
        return False
    levs=[]
    for line in fin:
        if line[0] == '#': continue
        tmp=line.split()
        for t in tmp:
            try:
                tt=int(t)
            except:
                return False
            if tt not in levs: levs.append(tt)
    return levs



# *****************************************************************************
# main program
ul=[]
ul.append("usage: %prog [options] filename [indices]")
ul.append("")
ul.append("Extract mixing percentages from an AUTOSTRUCTURE olg file.  There")
ul.append("are three modes of operation: (1) full, show mixing percentages for")
ul.append("individual levels, (2) conf, sum percentages over configuration")
ul.append("labels, and (3) term, sum percentages over term labels.")
usage=""
for u in ul: usage+=u+'\n'

parser=OptionParser(usage=usage)
parser.add_option("-m","--mode",dest="mode",default="full",type='string',help="select operation mode: full (default), conf, term",metavar="MODE")
parser.add_option("-t","--tol",dest="tol",default="2",type='int',help="specify number of decimal points to print (default: 2)",metavar="TOL")
parser.add_option("-f","--filter",dest="filter",default=None,type='string',help="file containing level indices",metavar="FILE")
parser.add_option("-o","--output",dest="output",default=None,type='string',help="write to file instead of screen",metavar="FILE")
parser.add_option("--terms",action="store_true",dest="terms",default=False,help="operate on LS terms in file instead of levels")
parser.add_option("--confmatrix",action="store_true",dest="confmat",default=False,help="write out configuration mixing matrix")

(options,args)=parser.parse_args()

# get input file name and check type
if len(args) < 1:
    parser.print_help()
    sys.exit(0)
else:
    filename=args[0]
    if not os.path.isfile(filename):
        print '*** Error - '+str(filename)+' not found ***'
        sys.exit(0)
    elif rlf_filetype(filename) != 'olg':
        print '*** Error - '+str(filename)+' is not an olg file ***'
        sys.exit(0)

# create level and mixing objects
levels=False
if not options.terms:
  try:
    levels=rap_states_lsj(filename)
    cup='LSJ'
  except:
    pass
if levels is False:
  try:
    levels=rap_states_ls(filename)
    cup='LS'
  except:
    print '*** Error - could not get energy data from olg file ***'
    sys.exit()
mixing=rap_mix(filename,levels)

# check for tolerance
tol=2
if options.tol:
    try:
        ttol=int(options.tol)
    except:
        print '*** Error - invalid value for the tolerance ***'
        sys.exit(0)
    tol=ttol

# check for output file
outfile=None
if options.output: 
    outfile=str(options.output)
    try:
        fout=open(outfile,'w')
    except:
        print '*** Error - could not open output file ***'
        sys.exit(0)
    else:
        fout.close()

# check if making configuration mixing matrix
if options.confmat:
    conf=levels.configurations()
    nconf=len(conf)
    cmap=[]
    for i in levels:
        cmap.append(conf.index(levels.configuration(i))+1)
    factor=[]
    for c in cmap: factor.append(1.0/float(cmap.count(c)))
    mixmat=mixing.matrix()
    smmat=[[0.0 for i in range(nconf)] for j in range(nconf)]
    for i in levels:
        ic=cmap[i-1]
        for j in levels:
            jc=cmap[j-1]
            smmat[ic-1][jc-1]+=mixmat[i-1][j-1]**2*factor[i-1]
    # print out matrix
    lines=[]
    for i in range(1,nconf+1): lines.append('# '+'%-3i'%i+': '+conf[i-1])
    lines.append('#')
    iform='%'+str(3)+'i'
    fform='%'+str(tol+2)+'.'+str(tol)+'f'
    line='#'+' '*(3)+'  '
    for i in range(1,nconf+1): line+=str(i).center(tol+2)+'  '
    lines.append(line)
    for i in range(1,nconf+1):
        line=iform%i+'  '
        for j in range(1,nconf+1): line+=fform%smmat[i-1][j-1]+'  '
        lines.append(line)
    if outfile:
        fout=open(outfile,'w')
        for line in lines: fout.write(line+'\n')
        fout.close()
    else:
        for line in lines: print line
else:   # regular operation
    # check for level indices
    levs=[]
    if options.filter != None:
        levs=read_levlist(options.filter)
        if not levs:
            print '*** Error - trouble reading level list from '+str(options.filter)+' ***'
            sys.exit(0)
    for a in args[1:]:
        try:
            l=int(a)
        except:
            print '*** Error - invalid level index given ***'
            sys.exit(0)
        if l not in levs: levs.append(l)
    levs.sort()
    if not levs: 
        print '*** Error - no level indices given ***'
        sys.exit(0)
    if levs[0] < 1 or levs[-1] > len(levels):
        print '*** Error - level index out of range ***'
        sys.exit(0)
    
    # check for mode
    mode='full'
    if options.mode:
        o=options.mode.lower()[:4]
        if o in ['full','conf','term']:
            mode=o
        else:
            print '*** Error - invalid operation mode given ***'
            sys.exit(0)
    
    # print out mixing percentages for each level
    for l in levs:
        coeffs=mixing.coeff(l)
        if not coeffs:
            print '*** Error - mixing coefficients could not be retrieved ***'
            sys.exit(0)
        line1=str(l)+": "+levels.label(l)
        line2='-'*len(line1)
        if outfile:
            try:
                fout=open(outfile,'a')
            except:
                print '*** Error - could not open output file ***'
                sys.exit(0)
            fout.write('\n\n')
            fout.write(line1+'\n')
            fout.write(line2+'\n')
            fout.close()
        else:
            print ''
            print ''
            print line1
            print line2
        if mode == 'full':
            p=sum_full(levels,coeffs)
            label='Level'
        elif mode == 'conf':
            p=sum_config(levels,coeffs)
            label='Configuration'
        elif mode == 'term':
            p=sum_term(levels,coeffs)
            label='Term'
        else:
            print "*** Error - sanity check; shouldn't have gotten here ***"
            sys.exit(0)
        show_coeffs(p,label,tol=tol,filename=outfile)
